"""
Processing of qm9 dataset.
"""

from os.path import join
import pandas as pd
import numpy as np

from helixwrapper.datasets.inmemory_dataset import InMemoryDataset


def get_default_qm9_task_names():
    """Get that default freesolv task names and return measured expt"""
    return ['homo', 'lumo', 'gap']


def load_qm9_dataset(data_path, task_names=None):
    """
    tbd
    """
    if task_names is None:
        task_names = get_default_qm9_task_names()

    csv_file = join(data_path, 'raw/qm9.csv')
    input_df = pd.read_csv(csv_file, sep=',')
    smiles_list = input_df['smiles']
    labels = input_df[task_names]

    data_list = []
    for i in range(len(labels)):
        data = {
            'smiles': smiles_list[i],
            'label': labels.values[i],
        }
        data_list.append(data)
    dataset = InMemoryDataset(data_list)
    return dataset
    

def get_qm9_stat(data_path, task_names):
    """Return mean and std of labels"""
    csv_file = join(data_path, 'raw/qm9.csv')
    input_df = pd.read_csv(csv_file, sep=',')
    labels = input_df[task_names].values
    return {
        'mean': np.mean(labels, 0),
        'std': np.std(labels, 0),
        'N': len(labels),
    }
